#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
下载微表情识别模型
"""

import os
import sys
import requests
import torch
import logging
from tqdm import tqdm

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def download_file(url, destination):
    """
    从URL下载文件到目标路径
    Args:
        url: 下载URL
        destination: 目标文件路径
    """
    if os.path.exists(destination):
        logger.info(f"文件已存在: {destination}")
        return True
    
    os.makedirs(os.path.dirname(destination), exist_ok=True)
    
    try:
        logger.info(f"从{url}下载文件...")
        response = requests.get(url, stream=True)
        response.raise_for_status()
        
        total_size = int(response.headers.get('content-length', 0))
        block_size = 1024
        
        with open(destination, 'wb') as f, tqdm(
                desc=destination,
                total=total_size,
                unit='B',
                unit_scale=True,
                unit_divisor=1024,
            ) as bar:
            for data in response.iter_content(block_size):
                f.write(data)
                bar.update(len(data))
        
        logger.info(f"文件下载完成: {destination}")
        return True
    except Exception as e:
        logger.error(f"下载文件时出错: {str(e)}")
        if os.path.exists(destination):
            os.remove(destination)
        return False

def create_mock_swin3d_model():
    """
    创建一个模拟的Swin3D模型文件
    用于在无法下载真实预训练模型时使用
    """
    model_path = "models/swin3d_microexpr.pth"
    if os.path.exists(model_path):
        logger.info(f"模型文件已存在: {model_path}")
        return True
    
    try:
        logger.info("创建模拟的Swin3D模型文件...")
        
        # 定义模型参数
        embed_dim = 96
        depths = [2, 2, 6, 2]
        num_heads = [3, 6, 12, 24]
        num_classes = 7  # 7种微表情类别
        
        # 创建更完整的模型状态字典
        state_dict = {}
        
        # 补丁嵌入层
        state_dict['patch_embed.proj.weight'] = torch.randn(embed_dim, 3, 2, 4, 4)
        state_dict['patch_embed.proj.bias'] = torch.randn(embed_dim)
        state_dict['patch_embed.norm.weight'] = torch.ones(embed_dim)
        state_dict['patch_embed.norm.bias'] = torch.zeros(embed_dim)
        
        # 位置编码
        state_dict['absolute_pos_embed'] = torch.randn(1, 4*28*28, embed_dim)
        
        # 各个阶段的层
        curr_dim = embed_dim
        for i, (depth, num_head) in enumerate(zip(depths, num_heads)):
            for j in range(depth):
                # 注意力层
                state_dict[f'layers.{i}.blocks.{j}.norm1.weight'] = torch.ones(curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.norm1.bias'] = torch.zeros(curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.attn.qkv.weight'] = torch.randn(curr_dim * 3, curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.attn.qkv.bias'] = torch.randn(curr_dim * 3)
                state_dict[f'layers.{i}.blocks.{j}.attn.proj.weight'] = torch.randn(curr_dim, curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.attn.proj.bias'] = torch.randn(curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.attn.relative_position_bias_table'] = torch.randn(49, num_head)
                
                # MLP层
                state_dict[f'layers.{i}.blocks.{j}.norm2.weight'] = torch.ones(curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.norm2.bias'] = torch.zeros(curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.mlp.fc1.weight'] = torch.randn(curr_dim * 4, curr_dim)
                state_dict[f'layers.{i}.blocks.{j}.mlp.fc1.bias'] = torch.randn(curr_dim * 4)
                state_dict[f'layers.{i}.blocks.{j}.mlp.fc2.weight'] = torch.randn(curr_dim, curr_dim * 4)
                state_dict[f'layers.{i}.blocks.{j}.mlp.fc2.bias'] = torch.randn(curr_dim)
            
            # 下采样层
            if i < len(depths) - 1:
                state_dict[f'layers.{i}.downsample.norm.weight'] = torch.ones(curr_dim)
                state_dict[f'layers.{i}.downsample.norm.bias'] = torch.zeros(curr_dim)
                state_dict[f'layers.{i}.downsample.reduction.weight'] = torch.randn(curr_dim * 2, curr_dim, 1, 2, 2)
                curr_dim *= 2
        
        # 最终分类层
        state_dict['norm.weight'] = torch.ones(curr_dim)
        state_dict['norm.bias'] = torch.zeros(curr_dim)
        state_dict['head.weight'] = torch.randn(num_classes, curr_dim)
        state_dict['head.bias'] = torch.randn(num_classes)
        
        # 包装为检查点格式
        checkpoint = {
            'model_state_dict': state_dict,
            'config': {
                'model_type': 'SWIN3D',
                'embed_dim': embed_dim,
                'depths': depths,
                'num_heads': num_heads,
                'num_classes': num_classes,
                'patch_size': (2, 4, 4),
                'window_size': (8, 7, 7),
            },
            'epoch': 0,
            'optimizer_state_dict': {},
            'scheduler_state_dict': {}
        }
        
        # 保存模型
        torch.save(checkpoint, model_path)
        logger.info(f"模拟模型文件已创建: {model_path}")
        return True
    except Exception as e:
        logger.error(f"创建模拟模型失败: {str(e)}")
        if os.path.exists(model_path):
            os.remove(model_path)
        return False

def download_swin3d_model():
    """下载Swin Transformer 3D模型"""
    model_path = "models/swin3d_microexpr.pth"
    if os.path.exists(model_path):
        logger.info(f"模型文件已存在: {model_path}")
        return True
    
    # 尝试从多个源下载
    sources = [
        # 微软官方模型
        "https://github.com/SwinTransformer/storage/releases/download/v1.0.4/swin_tiny_patch244_window877_kinetics400_1k.pth",
        # 备用源
        "https://download.openmmlab.com/mmaction/v1.0/recognition/swin/swin_tiny_patch244_window877_kinetics400_1k.pth"
    ]
    
    for url in sources:
        try:
            logger.info(f"尝试从 {url} 下载Swin3D模型...")
            if download_file(url, model_path):
                logger.info("Swin3D模型下载成功")
                return True
        except Exception as e:
            logger.error(f"从 {url} 下载失败: {str(e)}")
    
    # 如果所有下载源都失败，创建一个模拟模型
    logger.warning("所有下载源都失败，将创建一个模拟模型")
    return create_mock_swin3d_model()

def download_face_detector_model():
    """下载人脸检测模型"""
    # 根据配置，项目使用OpenCV级联分类器进行人脸检测
    # 实际上OpenCV级联分类器模型已经内置在OpenCV中，无需额外下载
    
    model_dir = "models/face_detector"
    os.makedirs(model_dir, exist_ok=True)
    
    # 下载OpenCV的级联分类器模型
    haarcascade_file = os.path.join(model_dir, "haarcascade_frontalface_default.xml")
    if os.path.exists(haarcascade_file):
        logger.info(f"人脸检测级联分类器已存在: {haarcascade_file}")
        return True
    
    try:
        # 从GitHub下载级联分类器XML文件
        url = "https://raw.githubusercontent.com/opencv/opencv/master/data/haarcascades/haarcascade_frontalface_default.xml"
        logger.info(f"下载OpenCV级联分类器...")
        return download_file(url, haarcascade_file)
    except Exception as e:
        logger.error(f"下载人脸检测模型失败: {str(e)}")
        return False

def main():
    """主函数"""
    logger.info("开始下载预训练模型...")
    
    # 检查CUDA可用性
    cuda_available = torch.cuda.is_available()
    logger.info(f"CUDA 可用: {cuda_available}")
    
    # 下载Swin3D模型
    if download_swin3d_model():
        logger.info("Swin3D模型准备完成")
    else:
        logger.error("Swin3D模型准备失败")
    
    # 下载人脸检测模型
    if download_face_detector_model():
        logger.info("人脸检测模型下载成功")
    else:
        logger.error("人脸检测模型下载失败")
    
    logger.info("预训练模型下载任务完成")

if __name__ == "__main__":
    main() 